#pragma once

#include "kernel_operator.h"
#include "op_common.h"
#include <cstdint>

#ifndef __SHAPE_MAX_DIM__
#define __SHAPE_MAX_DIM__ 64
#endif

namespace AscendC {

class BroadcastHelper {
public:
  ShapeData broadcastShape;
  int64_t shapeMaxDim = 64;
  int64_t broadcastIndex[__SHAPE_MAX_DIM__];
  int64_t offset = 0;
  __aicore__ inline BroadcastHelper() {}

  __aicore__ inline void setIndexByOffset(int64_t value) {
    offset = value;
    int64_t size_temp = broadcastShape.size;
    for (int64_t j = 0; j < broadcastShape.dimNum; ++j) {
      size_temp = size_temp / broadcastShape[j];
      broadcastIndex[j] = value / size_temp;
      value = value % size_temp;
    }
  }

  __aicore__ inline void broadcastIndexIncrease(int64_t value) {
    offset += value;
    for (auto j = broadcastShape.dimNum - 1; j >= 0; --j) {
      auto newValue = broadcastIndex[j] + value;
      value = newValue / broadcastShape[j];
      newValue = newValue % broadcastShape[j];
      broadcastIndex[j] = newValue;
      if (value == 0) {
        break;
      }
    }
  }

  __aicore__ inline int64_t
  broadcastIndex2targetOffset(ShapeData &targetShape) {
    int64_t ret = 0;
    int64_t index = 0;
    int64_t size_temp = targetShape.size;
    for (auto i = 0; i < broadcastShape.dimNum; ++i) {
      size_temp = size_temp / targetShape[i];
      index = broadcastIndex[i];
      if (index >= targetShape[i]) {
        index = 0;
      }
      ret += index * size_temp;
    }
    return ret;
  }

  __aicore__ inline int64_t Index2Offset(int64_t *targetIndex) {
    int64_t ret = 1;
    for (auto i = 0; i < broadcastShape.dimNum; ++i) {
      ret *= targetIndex[i];
    }
    return ret;
  }

  __aicore__ inline void copyShapeIndex(int64_t *from, int64_t *out,
                                        int64_t size) {
    for (auto i = 0; i < size; ++i) {
      out[i] = from[i];
    }
  }

  template <typename T>
  __aicore__ inline void DataCopy(LocalTensor<T> dst, GlobalTensor<T> from,
                                  ShapeData &targetShape, int64_t calcCount) {
    constexpr int64_t alignSize = 32 / sizeof(T);
    auto finalInddex = broadcastShape.dimNum - 1;
    auto broadcastfinalDim = broadcastShape[finalInddex];
    int64_t dst_offset = 0;
    int64_t from_offset = 0;
    // dst.SetValue(0, T(123));
    for (int counter = calcCount; counter > 0;) {
      from_offset = broadcastIndex2targetOffset(targetShape);
      auto finalNeedCount =
          broadcastShape[finalInddex] - broadcastIndex[finalInddex];
      if (finalNeedCount > counter) {
        finalNeedCount = counter;
      }
      // 拷贝数据
      if (broadcastShape[finalInddex] == targetShape[finalInddex]) {
        myDataCopy(dst, from, finalNeedCount, dst_offset, from_offset);
      } else {
        constexpr uint32_t max_loop_count = 1024;
        auto target_value = from.GetValue(from_offset);
        for (auto i = 0; i < finalNeedCount; ++i) {
          dst.SetValue(dst_offset + i, target_value);
        }
      }
      counter -= finalNeedCount;
      dst_offset += finalNeedCount;
      broadcastIndexIncrease(finalNeedCount);
    }
  }
};

} // namespace AscendC